import subprocess
import rrd_base
import re
import time

import color

try:
    from pyrrd.rrd import DataSource, RRA, RRD
    from pyrrd.graph import DEF, CDEF, VDEF
    from pyrrd.graph import LINE, AREA, GPRINT
    from pyrrd.graph import ColorAttributes, Graph
except:
    print "Couldn't load required modules. You probably need to install it with:"
    sys.exit('sudo apt-get install python-pyrrd python-pysnmp4')

class RrdTasks(rrd_base.RrdBase):
    def __init__(self, domain, top_lines):
        super(RrdTasks,self).__init__(domain=domain, group="cpu", instance="tasks")
        self.top_lines = top_lines

    def create_file(self):
        dataSources = []
        for ds in ['running', 'sleeping', 'stopped', 'zombie']:
            dataSources.append(DataSource( dsName=ds, dsType='GAUGE', heartbeat=self.heartbeat()))

        roundRobinArchives = []
        for func in ['AVERAGE', 'MAX', 'MIN']:
            for step,row in self.rra_pattern():
                roundRobinArchives.append(RRA(cf=func, xff=0.5, steps=step, rows=row))

        self.myRRD = RRD(self.filename, ds=dataSources, rra=roundRobinArchives, step=self.data_step())
        self.myRRD.create()

    def update(self):
        for line in self.top_lines:
            if re.match("Tasks:", line):
                line = re.sub("[^0-9\. ]", "", line).split()
                # ['243', '1', '239', '1', '2']
                self.myRRD.bufferValue(int(round(time.time())),
                                       float(line[1]),
                                       float(line[2]),
                                       float(line[3]),
                                       float(line[4]))
                self.myRRD.update()
                break

    def graph(self):
        def_ru    = DEF(rrdfile=self.myRRD.filename, vname='running',  cdef="AVERAGE", dsName='running')
        def_sl    = DEF(rrdfile=self.myRRD.filename, vname='sleeping', cdef="AVERAGE", dsName='sleeping')
        def_st    = DEF(rrdfile=self.myRRD.filename, vname='stopped',  cdef="AVERAGE", dsName='stopped')
        def_zo    = DEF(rrdfile=self.myRRD.filename, vname='zombie',   cdef="AVERAGE", dsName='zombie')
        area_ru   = AREA(defObj=def_ru, color=color.RGBA['navy'], legend="Running", stack=True)
        area_sl   = AREA(defObj=def_sl, color=color.RGBA['green'],legend="Sleeping",stack=True)
        area_st   = AREA(defObj=def_st, color=color.RGBA['blue'], legend="Stopped", stack=True)
        area_zo   = AREA(defObj=def_zo, color=color.RGBA['red'],  legend="Zombie",  stack=True)

        for tf in self.graph_time_frames().keys():
            g = Graph(self.image_name(tf),
                      start=int(time.time())-self.graph_time_frames()[tf], end=int(time.time()),
                      lower_limit=0.000001,
                      title='"Number of Tasks for {0} - Updated {1}"'.format(self.data_domain, time.strftime("%Y%m%d at %H:%M")),
                      height=self.graph_height, width=self.graph_width,
                      x_grid=self.grid[tf],
                      step=self.graph_step_times[tf])
            g.data.extend([ def_ru, def_sl, def_st, def_zo, area_ru, area_sl, area_st, area_zo ])
            g.write()

            if False:
                g = Graph(self.icon_name(tf),
                          start=int(time.time())-self.graph_time_frames()[tf], end=int(time.time()),
                          lower_limit=0.000001,
                          no_legend=True,
                          height=self.icon_height, width=self.icon_width,
                          step=self.graph_step_times[tf])
                g.data.extend([ def_ru, def_sl, def_st, def_zo, area_ru, area_sl, area_st, area_zo ])
                g.write()
